import os
import sys

import cv2
import numpy as np
import selectivesearch


sys.path.append('..')


def filter_out(src_frame):
    '''
    只提取黑色部分
    :param src_frame:
    :return:
    '''
    if src_frame is not None:
        hsv = cv2.cvtColor(src_frame, cv2.COLOR_BGR2HSV)
        lower_val = np.array([0, 0, 0])
        upper_val = np.array([179, 255, 127])
        mask = cv2.inRange(hsv, lower_val, upper_val)
        mask_inv = cv2.bitwise_not(mask)
        return mask_inv


files = os.listdir("./raw_data")
for f in files:
    print(f)
    img = cv2.imread("./raw_data/" + f)#读这张图片
    cv2.imshow("source", img)#展示原图

    cv2.threshold(img, 150, 255, cv2.THRESH_BINARY, img)  # 二值化
    t_width = img.shape[1]
    t_height = img.shape[0]

    cv2.imshow("after threshold", img)

    black_img = filter_out(img.copy())  # 提取黑色像素

    cv2.imshow("get black", black_img)

    kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (3, 3))
    black_img = cv2.erode(black_img, kernel)  # 腐蚀
    cv2.imshow("get erode", black_img)  # 腐蚀
    cv2.imwrite("temp.png", black_img)
    bgrimg = cv2.imread("temp.png")  # 重新存储一下是为了避开python版本的bug
    # cv2.imshow("reread", bgrimg)

    black_img = bgrimg
    useOpencv = True

    if (useOpencv):
        img_cpy = img.copy()
        img = cv2.cvtColor(black_img, cv2.COLOR_BGR2GRAY)  # rgb 变成 灰度图，才能用于 findContours 函数
        contours, hierarchy = cv2.findContours(img, cv2.RETR_LIST, cv2.CHAIN_APPROX_SIMPLE)  # 寻找联通区域

        for con in contours:
            box = cv2.boundingRect(con)  # x，y，w，h 用外接矩形框框住联通区域
            if box[2] < t_width and (box[3] > t_height * 1.0 / 2):  # 保留候选区域，符合这个条件的，就有可能是验证码区域
                # 开始的y坐标:结束的y坐标,开始x:结束的x
                cv2.rectangle(img_cpy, (box[0], box[1]), (box[0] + box[2], box[1] + box[3]), (0, 0, 255),
                              thickness=1)  # 绘制候选区域
                saved = black_img[box[1]:box[1] + box[3], box[0]:box[0] + box[2]]  # 裁剪出来
                # cv2.imwrite("./split_data/" + str(getRandomInt()) + ".png", saved)  # 保存好
                # 只要能正确的裁剪出目标区域，基本确定这个验证码就非常容易破解了
                # 怎么裁剪，需要做实验根据实际情况确定

        cv2.imshow("result", img_cpy)
        cv2.waitKey(1000)


    else:  # 这部分是使用selctive 算法，这个算法暂时不研究，有opencv 用这个函数的意义不大
        # black_img.shape[0]=3
        # print("shape.channes",black_img.shape[2])

        img_lbl, regions = selectivesearch.selective_search(bgrimg, scale=500, sigma=0.9, min_size=10)
        t_width = img.shape[1]
        t_height = img.shape[0]
        print("width:", str(t_width), " height:" + str(t_height))
        for reg in regions:
            rects = reg['rect']
            width = rects[2]
            height = rects[3]
            print("target width:" + str(width) + " targetHeight:" + str(height))
            if height > t_height / 2 and width < t_width / 4:
                cv2.rectangle(img, (rects[0], rects[1]), (rects[0] + rects[2], rects[1] + rects[3]), (0, 0, 255),
                              thickness=1)
        cv2.imshow("boxed ", img)
        cv2.waitKey(1000)
